# Grid trading strategies
# 
# Author: Mario
# 
# Version: v.20200405


import copy
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from pathlib import Path
from download import get_history_price_range
import logging
logging.getLogger("grid").setLevel(logging.INFO)

# only support tick data
FILENAME = "tushare_512000.csv"
DATA_PATH = Path("etc") / FILENAME
DATA = pd.read_csv(DATA_PATH)
print(get_history_price_range(DATA))
# Pre-setting: capital, premium (10^-3RMB), grid_bottom(the lowest price)
CAPITAL = 100 * (10 ** 4)
PREMIUM = 6 * (10 ** (-3))
# TODO: the higher bottom, the higher premium. But the grid size will be squeezed if the bottom is high.
GRID_BOTTOM = 70 * (10 ** (-3))
GRID_DISTRIBUTION = np.array([])
TRADE_FEE = .00035


def init_grid_setting(set_grid_size=None):
    max_price, min_price = get_history_price_range(DATA)
    if min_price > GRID_BOTTOM:
        min_price = GRID_BOTTOM

    grid_total = int((max_price - min_price) // 10 ** (-3))
    # the least size is 100 shares
    global GRID_DISTRIBUTION
    GRID_DISTRIBUTION = np.arange(int(min_price * 10 ** 3), int(max_price * 10 ** 3) + 1) / 10 ** 3

    grid_sequence = GRID_DISTRIBUTION * 10 ** 2
    grid_max = grid_sequence.cumsum()[-1]
    grid_size = (CAPITAL // grid_max) * 100 if not set_grid_size else set_grid_size
    assert grid_size > 0, "error: your capital is too small to trade."
    return dict(grid_size=grid_size, grid_total=grid_total,
                grid_max=round(max_price, 3), grid_min=round(min_price, 3))


# In-day:
def get_one_trade(date=None):
    # limit the trade window within [09:30:00, 15:00:00]
    min_trade_time, max_trade_time = "09:30:00", "15:00:00"
    date = date if date else DATA.date.unique()[0]
    sub_table = DATA[DATA.date == date]
    sub_table = sub_table[(sub_table.time >= min_trade_time) & (sub_table.time <= max_trade_time)]
    sub_table.index = range(len(sub_table))
    return sub_table


# set the account of the trader
class Trader:
    # Trader will be called in each period to ask if trade or keep.
    def __init__(self, **grid_setting):
        self.capital = CAPITAL
        self.base_price = 0
        self.grid_size = grid_setting.get("grid_size")
        self.grid_min = grid_setting.get("grid_min")
        self.grid_max = grid_setting.get("grid_max")
        self.grid_premium = PREMIUM
        # NB. at initial stage, can_trade is True
        # can_trade is True when a near sell is made and the base_price is changed,
        # can_trade is False when the price is lower than minimum and the price is filled
        self.can_trade = True
        # NB. account records the share owned in the format:
        """
        {
            date1: {
                price1: quantity1,
                price2: quantity2
            },
            date2: {
                price1: quantity1,
                ...
            }
        }
        """
        # Trade Convention:
        # set a base price like: 0.900, and buy if the price hits 0.899,
        # and sell if the price hits 0.904 when the previous transaction 0.899 is valid(cross a day)
        self.account = pd.DataFrame(GRID_DISTRIBUTION, columns=["price"])
        self.account["quantity"] = 0
        # price-buy_date pair is unique because only when the previous transaction is released,
        # new transaction can be thus made.
        self.account["buy_date"] = ""

        # NB. history records all the transactions in the format:
        # 1. bill_date(with time), 2. buy_date(with time) (if sell), 3. mode(sell, buy),
        # 4. price, 5. quantity, 6. revenue
        self.history = []

        # NB. dynamic revenue records the dynamic revenue logging for each trading day
        self.dynamic_revenue = {}
        # NB. daily revenue records the revenue per trading day
        self.daily_revenue = {}
        # NB. holding percent records the percentage of share holdings per trading day
        self.holding_percent = {}

    @staticmethod
    def generate_record(transaction):
        record = dict(bill_date="", buy_date="", mode="",
                      price=0, quantity=0, revenue=0)
        for key, value in transaction.items():
            record[key] = value
        return record

    @staticmethod
    def generate_grids_between_prices(low, high):
        return list(np.array(range(int(low * 10 ** 3), int(high * 10 ** 3) + 1)) / 10 ** 3)

    @staticmethod
    def generate_transaction():
        return dict(price=0, quantity=0, buy_date="", mode="")

    def record_in(self, record: dict):
        self.history += [record]
        self.log(record)

    def buffer(self, transactions):
        # TODO: buffer helps manage the transactions that cannot be proceeded at once.
        pass

    def abnormal_schedule(self):
        # NB. abnormal schedule catches abnormal downwards or upwards according to the account status,
        # downwards: buy grids with .generate_grids_between_prices
        # upwards: sell grids with .generate_grids_between_prices
        # TODO: this method can only be called when using the minute/second K-information.
        pass

    def schedule(self, now, price, quantity):
        # :return: {price, quantity, buy_date}
        transactions = []
        if price > self.grid_max:
            return transactions

        if price < self.grid_min:
            # in this case, buy in all the grids between the minimum and the price
            prices = self.generate_grids_between_prices(price, self.grid_min)
            for pr in prices:
                transaction = copy.deepcopy(self.generate_transaction())
                transaction["price"] = pr
                transaction["quantity"] = self.grid_size
                transaction["buy_date"] = now
                transaction["mode"] = "buy"
                transactions.append(transaction)
            return transactions

        # sell request: should scan all the account's entries with quantity > 0
        sell_permit_price = round(price - self.grid_premium, 3)
        sell_permit = (self.account.price == sell_permit_price) \
                      & (self.account.buy_date < now) \
                      & (self.account.quantity > 0)
        record = self.account[sell_permit]
        if not record.empty:
            mode = "sell"
            # TODO: if the policy is strict, sell one grid otherwise can sell all (here sell one).
            record = self.account[sell_permit].reset_index().drop(columns="index")
            transaction = record.iloc[0].to_dict()
            if transaction["quantity"] > quantity:
                transaction["quantity"] = quantity
            transaction["bill_date"] = now
            transaction["mode"] = mode
            transactions.append(transaction)
            self.can_trade = True

        # NB. as the grid_max price is set, the empty error won't happen
        # but to prevent it, check the empty array:
        buy_quota = self.account[self.account.price == price].quantity.values
        if buy_quota.size != 0:
            if buy_quota[0] == 0:
                mode = "buy"
                transaction = copy.deepcopy(self.generate_transaction())
                transaction["price"] = price
                transaction["quantity"] = self.grid_size if self.grid_size <= quantity else quantity
                transaction["buy_date"] = now
                transaction["mode"] = mode
                transactions.append(transaction)

        return transactions

    @staticmethod
    def log(record):
        # log the account status and transactions
        logging.info("[{}] {}: {} at {} ".format(record["bill_date"], record["mode"],
                                                record["quantity"], record["price"]))

    def revenue(self):
        profit = 0
        if self.history:
            for each in self.history:
                profit += each.get("revenue", 0)
        return profit

    def statistics(self):
        stat = pd.DataFrame()
        if self.history:
            for each in self.history:
                line = pd.DataFrame(each, index=[0])
                stat = pd.concat([stat, line], axis=0)
            stat.index = range(len(stat))
        return stat

    @staticmethod
    def transaction_fee(request, fee_rate=TRADE_FEE):
        return request * fee_rate

    def buy(self, transaction):
        # check the status of capital
        buy_request = transaction["price"] * transaction["quantity"]
        if self.capital < buy_request:
            logging.info("[Error] capital: %.2f, buy: %.2f" % (self.capital, buy_request))
            return False
        # update the account
        index = self.account[self.account.price == transaction["price"]].index[0]
        # pandas warning: use .loc[row_indexer,col_indexer] = value instead
        self.account.quantity.iloc[index] = transaction["quantity"]
        self.account.buy_date.iloc[index] = transaction["buy_date"]
        self.base_price = transaction["price"]
        # record in the history
        record = self.generate_record(transaction)
        record["revenue"] = -self.transaction_fee(buy_request)
        record["bill_date"] = transaction["buy_date"]
        self.record_in(record)
        self.capital += -(buy_request + self.transaction_fee(buy_request))
        return True

    def sell(self, transaction):
        index = self.account[self.account.price == transaction["price"]].index[0]
        sell_quantity = self.account.quantity.iloc[index]
        if sell_quantity > transaction["quantity"]:
            self.account.quantity.iloc[index] = sell_quantity - transaction["quantity"]
        else:
            self.account.quantity.iloc[index] = 0
            transaction["quantity"] = sell_quantity

        sell_request = transaction["price"] * transaction["quantity"]
        record = self.generate_record(transaction)
        record["revenue"] = self.grid_premium * transaction["quantity"] - self.transaction_fee(sell_request)
        record["bill_date"] = transaction["buy_date"]
        self.record_in(record)
        # update the capital of account
        self.capital += sell_request - self.transaction_fee(sell_request)
        return True

    def update_dynamic_revenue(self, now, price):
        # update dynamic revenue only when the end of the trading day
        previous = self.account[self.account.buy_date < now]
        previous["dynamic"] = (previous.price - price) * previous.quantity
        dynamic = previous[["buy_date", "dynamic"]].groupby(["buy_date"]).agg(sum).reset_index()
        # log for each trading day (curve)
        self.dynamic_revenue[now] = round(dynamic.dynamic.sum(), 3)

    def update_daily_revenue(self, now):
        self.daily_revenue[now] = round(self.revenue(), 3)

    def update_holding_percent(self, now):
        # TODO: update the shareholdings percent for each trading date. (error)
        self.holding_percent[now] = (self.account.price * self.account.quantity).sum() / CAPITAL

    def trade(self, now, price, quantity, signal):
        # initialize when the base_price is zero
        if self.base_price == 0:
            self.base_price = price

        if signal == "wait":
            return

        # action is a list: allow the trader to operate at the same time
        actions = self.schedule(now, price, quantity)
        if actions:
            for action in actions:
                if action.get("mode") == "buy" and self.can_trade and signal == "buy":
                    self.can_trade = self.buy(action)
                elif action.get("mode") == "sell" and self.can_trade and signal == "sell":
                    self.can_trade = self.sell(action)
                else:
                    logging.info("[%s] capital: %.2f, revenue: %.2f" % (now, self.capital, self.revenue()))


# test on a one-time trade in a single trading day:
def trade_once(trader, table):
    # NB. buy-bid indicates: `selling` while sell-bid indicates: `buying`.
    for index in table.index:
        bid = table.iloc[index].to_dict()
        now = bid["date"]
        price = round(bid["price"], 3)
        # unit calibration: shares should be enlarged 100 times.
        quantity = bid["volume"] * 100
        # request for the trader (Trader)
        signal = "wait"
        if bid["type"] == "卖盘":
            signal = "buy"
        elif bid["type"] == "买盘":
            signal = "sell"
        trader.trade(now, price, quantity, signal)
    # NB. the price is the `close price` of the day(now)
    trader.update_dynamic_revenue(now, price)
    trader.update_daily_revenue(now)
    trader.update_holding_percent(now)
    return trader


def revenue_cross_validate(trader):
    x = list(trader.__getattribute__("daily_revenue").keys())
    daily_revenue = list(trader.__getattribute__("daily_revenue").values())
    dynamic_revenue = list(trader.__getattribute__("dynamic_revenue").values())
    fig, ax1 = plt.subplots(figsize=(10, 6), dpi=100)
    # plot both of the revenues and show the `danger zone` where the net profit is negative.
    ax1.bar(x, daily_revenue, color="orange", alpha=0.7, label="Daily Revenue")
    ax1.bar(x, dynamic_revenue, color="skyblue", alpha=0.7, label="Dynamic Revenue")
    ax1.set_ylabel("Revenue (RMB)")
    plt.legend(loc=2)
    # show the accumulated return rates:
    holdings = list(trader.__getattribute__("holding_percent").values())
    # transform the daily return in % unit (x100).
    daily_return = (np.array(daily_revenue) + np.array(dynamic_revenue)) / (np.array(holdings) * CAPITAL) * 100
    ax2 = ax1.twinx()
    ax2.plot(x, daily_return, color="red", marker="o", alpha=0.5, label="Daily Return")
    ax2.plot(x, np.zeros(daily_return.size), color="grey", alpha=0.6)
    plt.title(str(DATA_PATH).replace(DATA_PATH.suffix, "",))
    ax2.set_ylabel("Return / Holding Percent (%)")
    plt.xlabel("Trade Days")
    plt.margins(0.01)
    plt.legend(loc=3)
    ax1.set_xticks([x[i] for i in range(1, len(x) + 1, 5)])
    ax1.set_xticklabels([x[i] for i in range(1, len(x) + 1, 5)], rotation="90")
    fig_save_path = (Path() / FILENAME.split(".")[0]).with_suffix(".png")
    fig.savefig(str(fig_save_path), format="png")
    print("find revenue plot at: [ %s ]" % str(fig_save_path.resolve()))


if __name__ == "__main__":
    # customise your own grid_size which can significantly alter your expected revenues.
    setting = init_grid_setting(set_grid_size=6000)
    trader = Trader(**setting)
    for date in DATA.date.unique()[: 100]:
        table = get_one_trade(date)
        trader = trade_once(trader, table)
        # update and log the account
        print("[%s] capital: %.3f, revenue: %.3f, float: %.3f" % (date, trader.capital, trader.revenue(),
                                                                  trader.dynamic_revenue[date]))

    # show the performance:
    revenue_cross_validate(trader)
